from matplotlib import colors as mcolors
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.collections import PolyCollection
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import json

def cal_pass():
    df = pd.read_csv('E:/data.csv')
    df = df[df['TRAVELTIME']<df['TRAVELTIME'].max()*0.8]
    df = df.groupby(['TRAVELTIME','HOUR'])['COUNT'].sum()
    df = df.reset_index()
    lists = df.values.tolist()
    lists.insert(0,['TRAVELTIME','HOUR','COUNT'])
    jsObj = json.dumps(lists)

    fileObject = open('E:/jsonFile.json', 'w')
    fileObject.write(jsObj)
    fileObject.close()
    # draw(df)
    print ('ok')

def draw(df):

    x = df['HOUR']
    y = df['TRAVELTIME']
    z = df['COUNT']
    # 生成画布、3D图形对象、三维散点图
    fig = plt.figure()
    ax = Axes3D(fig)
    ax.scatter(x,y,z)

    # 设置坐标轴显示以及旋转角度
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    ax.view_init(elev=10, azim=235)
    plt.show()

def t(df):
    colors = [mcolors.to_rgba(c, alpha=0.6)
              for c in plt.rcParams['axes.prop_cycle'].by_key()['color']]

    poly = PolyCollection([df['TRAVELTIME'].tolist(),df['COUNT'].tolist()], facecolors=colors)
    poly.set_color(colors)
    poly.set_alpha(0.7)

    fig = plt.figure()
    plt.rcParams['font.sans-serif'] = ['simhei']
    plt.rcParams['axes.unicode_minus'] = False
    ax = fig.gca(projection='3d')
    ax.add_collection3d(poly, zs=df['COUNT'], zdir='y')

    ax.set_xlabel('X-行驶时间（秒）')
    ax.set_xlim3d(df['TRAVELTIME'].min(), df['TRAVELTIME'].max())
    ax.set_ylabel('Y-时间（小时数）')
    ax.set_ylim3d(0, 24)
    ax.set_zlabel('Z-行驶时间次数')
    ax.set_zlim3d(df['COUNT'].min(), df['COUNT'].max())

    ax.set_title("行驶时间频谱图")

    plt.show()


if __name__ == '__main__':
    cal_pass()